#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Excel处理器模块
负责Excel文件的读取和写入操作
"""

import pandas as pd
import logging
from pathlib import Path
import os
from typing import Dict, List, Tuple, Optional, Any, Union
from openpyxl import load_workbook
from openpyxl.utils import get_column_letter
from openpyxl.styles import Alignment
from openpyxl.worksheet.table import Table, TableStyleInfo

from utils.logger import get_logger

# 获取日志记录器
logger = get_logger(__name__)

class ExcelProcessor:
    """Excel文件处理器类"""
    
    def __init__(self):
        """初始化Excel处理器"""
        # 设置pandas显示选项
        pd.set_option('display.max_columns', None)
        pd.set_option('display.width', None)
    
    def read_excel(self, file_path, sheet_name=0, header=0, skiprows=None, nrows=None):
        """
        读取Excel文件
        
        Args:
            file_path (str): Excel文件路径
            sheet_name (str or int, optional): 工作表名称或索引，默认为0（第一个工作表）
            header (int or list of ints, optional): 作为列名的行，默认为0
            skiprows (list-like, optional): 要跳过的行索引
            nrows (int, optional): 要读取的行数
            
        Returns:
            pandas.DataFrame: 读取的数据框，如果读取失败则返回None
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                logger.error(f"文件不存在: {file_path}")
                return None
            
            # 检查文件扩展名
            file_ext = Path(file_path).suffix.lower()
            if file_ext not in ['.xlsx', '.xls']:
                logger.error(f"不支持的文件格式: {file_ext}")
                return None
            
            logger.info(f"正在读取Excel文件: {file_path}, 工作表: {sheet_name}")
            
            # 读取Excel文件
            df = pd.read_excel(
                file_path,
                sheet_name=sheet_name,
                header=header,
                skiprows=skiprows,
                nrows=nrows,
                engine='openpyxl' if file_ext == '.xlsx' else 'xlrd'
            )
            
            logger.info(f"成功读取Excel文件，共 {len(df)} 行，{len(df.columns)} 列")
            return df
            
        except Exception as e:
            logger.error(f"读取Excel文件失败: {str(e)}")
            return None
    
    def get_excel_structure(self, file_path: str) -> Dict[str, Any]:
        """
        获取Excel文件的完整结构信息，包括合并单元格等
        
        Args:
            file_path (str): Excel文件路径
            
        Returns:
            Dict: 包含Excel结构信息的字典，格式如下：
            {
                'sheets': {
                    'sheet_name1': {
                        'dimensions': 'A1:Z100',
                        'merged_cells': [
                            {'min_col': 1, 'min_row': 1, 'max_col': 2, 'max_row': 3, 'value': '合并单元格值'}
                        ],
                        'tables': [],
                        'total_rows': 100,
                        'total_cols': 26
                    }
                }
            }
        """
        try:
            if not os.path.exists(file_path):
                logger.error(f"文件不存在: {file_path}")
                return None
            
            file_ext = Path(file_path).suffix.lower()
            if file_ext != '.xlsx':
                logger.error("获取Excel结构仅支持.xlsx格式文件")
                return None
            
            logger.info(f"正在分析Excel文件结构: {file_path}")
            
            # 使用openpyxl加载工作簿
            wb = load_workbook(file_path, read_only=True)
            structure = {'sheets': {}}
            
            for sheet_name in wb.sheetnames:
                ws = wb[sheet_name]
                sheet_info = {
                    'dimensions': ws.dimensions,
                    'merged_cells': [],
                    'tables': [],
                    'total_rows': ws.max_row,
                    'total_cols': ws.max_column
                }
                
                # 获取合并单元格信息
                for merged_cell in ws.merged_cells.ranges:
                    min_col, min_row, max_col, max_row = merged_cell.bounds
                    # 获取合并单元格的值（左上角单元格的值）
                    value = ws.cell(row=min_row, column=min_col).value
                    sheet_info['merged_cells'].append({
                        'min_col': min_col,
                        'min_row': min_row,
                        'max_col': max_col,
                        'max_row': max_row,
                        'value': value,
                        'range': merged_cell.coord
                    })
                
                # 获取表格信息
                for table in ws.tables.values():
                    sheet_info['tables'].append({
                        'name': table.name,
                        'range': table.ref
                    })
                
                structure['sheets'][sheet_name] = sheet_info
            
            wb.close()
            logger.info(f"成功分析Excel文件结构，包含 {len(structure['sheets'])} 个工作表")
            return structure
            
        except Exception as e:
            logger.error(f"分析Excel文件结构失败: {str(e)}")
            return None
    
    def read_with_merged_info(self, file_path: str, sheet_name: Union[str, int] = 0) -> Tuple[Optional[pd.DataFrame], Optional[List[Dict]]]:
        """
        读取Excel文件并同时获取合并单元格信息
        
        Args:
            file_path (str): Excel文件路径
            sheet_name (str or int, optional): 工作表名称或索引，默认为0
            
        Returns:
            Tuple: (DataFrame, merged_cells_info)，如果失败则返回(None, None)
        """
        try:
            # 首先读取数据
            df = self.read_excel(file_path, sheet_name=sheet_name)
            if df is None:
                return None, None
            
            # 然后获取合并单元格信息
            structure = self.get_excel_structure(file_path)
            if structure is None:
                return df, None
            
            # 获取指定工作表的合并单元格信息
            if isinstance(sheet_name, int):
                sheet_names = list(structure['sheets'].keys())
                if sheet_name < len(sheet_names):
                    sheet_name = sheet_names[sheet_name]
                else:
                    return df, None
            
            merged_cells = structure['sheets'].get(sheet_name, {}).get('merged_cells', [])
            
            # 调整合并单元格信息，使其与DataFrame的索引对应
            # 注意：Excel行号从1开始，而DataFrame索引从0开始
            # 但由于pandas读取时header处理可能改变索引，这里返回原始Excel的合并单元格信息
            logger.info(f"成功读取Excel文件并获取 {len(merged_cells)} 个合并单元格信息")
            return df, merged_cells
            
        except Exception as e:
            logger.error(f"读取Excel文件并获取合并单元格信息失败: {str(e)}")
            return None, None
    
    def write_excel(self, df, file_path, sheet_name='Sheet1', index=False, header=True,
                    freeze_panes=None, format_columns=None):
        """
        将数据写入Excel文件
        
        Args:
            df (pandas.DataFrame): 要写入的数据框
            file_path (str): 输出文件路径
            sheet_name (str, optional): 工作表名称，默认为'Sheet1'
            index (bool, optional): 是否写入索引列，默认为False
            header (bool, optional): 是否写入列名，默认为True
            freeze_panes (tuple, optional): 冻结窗格的位置，例如(1, 0)表示冻结第一行
            format_columns (dict, optional): 列格式化配置
            
        Returns:
            bool: 是否成功写入
        """
        try:
            # 确保输出目录存在
            output_dir = os.path.dirname(file_path)
            if output_dir:
                Path(output_dir).mkdir(parents=True, exist_ok=True)
            
            # 确保文件扩展名是.xlsx
            if not file_path.lower().endswith('.xlsx'):
                file_path += '.xlsx'
            
            logger.info(f"正在写入Excel文件: {file_path}, 工作表: {sheet_name}")
            
            # 创建ExcelWriter对象
            with pd.ExcelWriter(file_path, engine='openpyxl') as writer:
                # 将数据框写入Excel
                df.to_excel(writer, sheet_name=sheet_name, index=index, header=header)
                
                # 获取工作表对象进行高级操作
                worksheet = writer.sheets[sheet_name]
                
                # 设置冻结窗格
                if freeze_panes:
                    worksheet.freeze_panes = freeze_panes
                
                # 自动调整列宽
                for column in worksheet.columns:
                    max_length = 0
                    column_letter = column[0].column_letter
                    for cell in column:
                        try:
                            if cell.value:
                                max_length = max(max_length, len(str(cell.value)))
                        except:
                            pass
                    adjusted_width = min(max_length + 2, 50)  # 限制最大宽度为50
                    worksheet.column_dimensions[column_letter].width = adjusted_width
            
            logger.info(f"成功写入Excel文件，共 {len(df)} 行，{len(df.columns)} 列")
            return True
            
        except Exception as e:
            logger.error(f"写入Excel文件失败: {str(e)}")
            return False
    
    def insert_row_after(self, file_path: str, sheet_name: Union[str, int] = 0, 
                        row_index: int = 0, data: Dict[str, Any] = None) -> bool:
        """
        在指定行后插入新行
        
        Args:
            file_path (str): Excel文件路径
            sheet_name (str or int, optional): 工作表名称或索引，默认为0
            row_index (int): 要在其后插入行的索引（Excel行号，从1开始）
            data (Dict): 要插入的数据，格式为{列名: 值}或{列索引: 值}
            
        Returns:
            bool: 是否成功插入
        """
        try:
            if not os.path.exists(file_path):
                logger.error(f"文件不存在: {file_path}")
                return False
            
            file_ext = Path(file_path).suffix.lower()
            if file_ext != '.xlsx':
                logger.error("插入行操作仅支持.xlsx格式文件")
                return False
            
            # 加载工作簿
            wb = load_workbook(file_path)
            
            # 获取工作表
            if isinstance(sheet_name, int):
                sheet_names = wb.sheetnames
                if sheet_name < len(sheet_names):
                    ws = wb[sheet_names[sheet_name]]
                else:
                    logger.error(f"工作表索引超出范围: {sheet_name}")
                    return False
            else:
                if sheet_name not in wb.sheetnames:
                    logger.error(f"工作表不存在: {sheet_name}")
                    return False
                ws = wb[sheet_name]
            
            # 插入行
            insert_position = row_index + 1  # 在指定行后插入
            ws.insert_rows(insert_position)
            
            # 填充数据（如果提供）
            if data:
                for key, value in data.items():
                    if isinstance(key, str):
                        # 假设第一行是表头，查找列索引
                        header_row = 1  # 默认第一行为表头
                        for col_idx in range(1, ws.max_column + 1):
                            if ws.cell(row=header_row, column=col_idx).value == key:
                                ws.cell(row=insert_position, column=col_idx).value = value
                                break
                    elif isinstance(key, int):
                        # 直接使用列索引
                        ws.cell(row=insert_position, column=key).value = value
            
            # 保存文件
            wb.save(file_path)
            wb.close()
            
            logger.info(f"成功在工作表'{sheet_name}'的行{row_index}后插入新行")
            return True
            
        except Exception as e:
            logger.error(f"插入行失败: {str(e)}")
            return False
    
    def create_group_structure(self, file_path: str, sheet_name: Union[str, int] = 0,
                              group_ranges: List[Dict[str, int]] = None) -> bool:
        """
        在Excel文件中创建组合结构（可折叠区域）
        
        Args:
            file_path (str): Excel文件路径
            sheet_name (str or int, optional): 工作表名称或索引，默认为0
            group_ranges (List): 要组合的行范围列表，格式为：
                [
                    {'start_row': 2, 'end_row': 5, 'outline_level': 1},
                    {'start_row': 6, 'end_row': 10, 'outline_level': 1}
                ]
            
        Returns:
            bool: 是否成功创建组合结构
        """
        try:
            if not os.path.exists(file_path):
                logger.error(f"文件不存在: {file_path}")
                return False
            
            file_ext = Path(file_path).suffix.lower()
            if file_ext != '.xlsx':
                logger.error("创建组合结构仅支持.xlsx格式文件")
                return False
            
            # 加载工作簿
            wb = load_workbook(file_path)
            
            # 获取工作表
            if isinstance(sheet_name, int):
                sheet_names = wb.sheetnames
                if sheet_name < len(sheet_names):
                    ws = wb[sheet_names[sheet_name]]
                else:
                    logger.error(f"工作表索引超出范围: {sheet_name}")
                    return False
            else:
                if sheet_name not in wb.sheetnames:
                    logger.error(f"工作表不存在: {sheet_name}")
                    return False
                ws = wb[sheet_name]
            
            # 创建组合结构
            if group_ranges:
                for group in group_ranges:
                    start_row = group.get('start_row')
                    end_row = group.get('end_row')
                    outline_level = group.get('outline_level', 1)
                    
                    if start_row and end_row:
                        # 设置行分组
                        for row_idx in range(start_row, end_row + 1):
                            ws.row_dimensions[row_idx].outline_level = outline_level
            
            # 保存文件
            wb.save(file_path)
            wb.close()
            
            logger.info(f"成功在工作表'{sheet_name}'创建组合结构")
            return True
            
        except Exception as e:
            logger.error(f"创建组合结构失败: {str(e)}")
            return False
    
    def get_sheet_names(self, file_path):
        """
        获取Excel文件中的所有工作表名称
        
        Args:
            file_path (str): Excel文件路径
            
        Returns:
            list: 工作表名称列表，如果失败则返回空列表
        """
        try:
            if not os.path.exists(file_path):
                logger.error(f"文件不存在: {file_path}")
                return []
            
            xl = pd.ExcelFile(file_path)
            sheet_names = xl.sheet_names
            logger.info(f"Excel文件包含工作表: {sheet_names}")
            return sheet_names
            
        except Exception as e:
            logger.error(f"获取工作表名称失败: {str(e)}")
            return []
    
    def merge_excel_files(self, input_files, output_file, sheet_name='Merged'):
        """
        合并多个Excel文件到一个文件中
        
        Args:
            input_files (list): 输入文件路径列表
            output_file (str): 输出文件路径
            sheet_name (str, optional): 合并后的工作表名称
            
        Returns:
            bool: 是否成功合并
        """
        try:
            logger.info(f"开始合并 {len(input_files)} 个Excel文件")
            
            # 合并所有文件的数据
            dfs = []
            for file_path in input_files:
                df = self.read_excel(file_path)
                if df is not None:
                    dfs.append(df)
            
            if not dfs:
                logger.error("没有成功读取任何文件")
                return False
            
            # 合并数据框
            merged_df = pd.concat(dfs, ignore_index=True)
            
            # 写入合并后的数据
            success = self.write_excel(merged_df, output_file, sheet_name=sheet_name)
            
            if success:
                logger.info(f"成功合并文件到 {output_file}，共 {len(merged_df)} 行数据")
            
            return success
            
        except Exception as e:
            logger.error(f"合并Excel文件失败: {str(e)}")
            return False
    
    def apply_merged_cells(self, file_path: str, sheet_name: Union[str, int] = 0,
                          merged_cells: List[Dict[str, int]] = None) -> bool:
        """
        在Excel文件中应用合并单元格设置
        
        Args:
            file_path (str): Excel文件路径
            sheet_name (str or int, optional): 工作表名称或索引，默认为0
            merged_cells (List): 要合并的单元格范围列表，格式为：
                [
                    {'min_col': 1, 'min_row': 1, 'max_col': 2, 'max_row': 3},
                    {'min_col': 4, 'min_row': 1, 'max_col': 5, 'max_row': 3}
                ]
            
        Returns:
            bool: 是否成功应用合并单元格设置
        """
        try:
            if not os.path.exists(file_path):
                logger.error(f"文件不存在: {file_path}")
                return False
            
            file_ext = Path(file_path).suffix.lower()
            if file_ext != '.xlsx':
                logger.error("应用合并单元格设置仅支持.xlsx格式文件")
                return False
            
            # 加载工作簿
            wb = load_workbook(file_path)
            
            # 获取工作表
            if isinstance(sheet_name, int):
                sheet_names = wb.sheetnames
                if sheet_name < len(sheet_names):
                    ws = wb[sheet_names[sheet_name]]
                else:
                    logger.error(f"工作表索引超出范围: {sheet_name}")
                    return False
            else:
                if sheet_name not in wb.sheetnames:
                    logger.error(f"工作表不存在: {sheet_name}")
                    return False
                ws = wb[sheet_name]
            
            # 应用合并单元格设置
            if merged_cells:
                for merge_info in merged_cells:
                    min_col = merge_info.get('min_col')
                    min_row = merge_info.get('min_row')
                    max_col = merge_info.get('max_col')
                    max_row = merge_info.get('max_row')
                    
                    if all([min_col, min_row, max_col, max_row]):
                        # 合并单元格
                        ws.merge_cells(
                            start_row=min_row,
                            start_column=min_col,
                            end_row=max_row,
                            end_column=max_col
                        )
                        
                        # 可以选择设置对齐方式
                        if 'alignment' in merge_info:
                            align_info = merge_info['alignment']
                            cell = ws.cell(row=min_row, column=min_col)
                            cell.alignment = Alignment(
                                horizontal=align_info.get('horizontal', 'general'),
                                vertical=align_info.get('vertical', 'bottom'),
                                wrap_text=align_info.get('wrap_text', False)
                            )
            
            # 保存文件
            wb.save(file_path)
            wb.close()
            
            logger.info(f"成功在工作表'{sheet_name}'应用 {len(merged_cells)} 处合并单元格设置")
            return True
            
        except Exception as e:
            logger.error(f"应用合并单元格设置失败: {str(e)}")
            return False